import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple, Dict
from cartopy.crs import EqualEarth

class CircularOrbit:
    """单个卫星圆轨道类"""
    
    def __init__(self, 
                 orbit_radius: float,                   # 轨道半径 >6371.0
                 inclination: float,                    # 轨道倾角 [0,pi]
                 raan: float = 0.0,                     # 升交点赤经 [0,2pi]
                 satellite_count: int = 2,              # 卫星数量
                 initial_true_anomaly: float = 0.0,     # 初始位置  [0,2pi]
                 color: str = 'r'):                     # 颜色参数
        self.orbit_radius = orbit_radius
        self.inclination = inclination
        self.raan = raan
        self.satellite_count = satellite_count
        self.initial_true_anomaly = initial_true_anomaly
        self.color = color
        
        self.EARTH_RADIUS = 6371.0
        if self.orbit_radius <= self.EARTH_RADIUS:
            raise ValueError(f"轨道半径必须大于地球半径 ({self.EARTH_RADIUS} km)")
            
        self.EARTH_MU = 3.986004418e5    # 地球引力常数
        self.orbital_period = 2 * np.pi * np.sqrt(self.orbit_radius**3 / self.EARTH_MU) # 计算轨道周期
        self._calculate_satellite_positions()
    
    def _calculate_satellite_positions(self) -> None:
        phase_step = 2 * np.pi / self.satellite_count                                   # 同轨道卫星之间的相位差
        self.satellite_true_anomalies = [
            self.initial_true_anomaly + i * phase_step for i in range(self.satellite_count) # 计算每个卫星的初始真近点角
        ]
    
    def get_satellite_position(self, satellite_index: int, time: float) -> Tuple[float, float, float]:
        """
        计算指定卫星在给定时间的位置
        
        参数:
            satellite_index (int): 卫星索引 (0到satellite_count-1)
            time (float): 时间点 (秒)
        
        返回:
            Tuple[float, float, float]: 卫星在ECEF坐标系中的位置 (x, y, z) (km)
        """
        # 检查卫星索引是否有效
        if satellite_index < 0 or satellite_index >= self.satellite_count:
            raise ValueError(f"卫星索引必须在0到{self.satellite_count-1}之间")
        
        # 计算卫星在轨道平面内的真近点角（考虑轨道运动）
        mean_motion = 2 * np.pi / self.orbital_period  # 平均运动 (rad/s)
        true_anomaly = self.satellite_true_anomalies[satellite_index] + mean_motion * time
        
        # 归一化到0-2π
        true_anomaly = true_anomaly % (2 * np.pi)
        '''
        # 计算轨道平面内的位置
        x_orbit = self.orbit_radius * np.cos(true_anomaly)
        y_orbit = self.orbit_radius * np.sin(true_anomaly)
        
        # 轨道平面到ECEF坐标系的转换
        # 1. 绕z轴旋转升交点赤经(RAAN)
        x_ecf1 = x_orbit * np.cos(self.raan) - y_orbit * np.sin(self.raan)
        y_ecf1 = x_orbit * np.sin(self.raan) + y_orbit * np.cos(self.raan)
        z_ecf1 = z_orbit  # 补充定义 z_ecf1
        
        # 2. 绕x轴旋转轨道倾角
        x_ecf = x_ecf1
        y_ecf = y_ecf1 * np.cos(self.inclination) - z_ecf1 * np.sin(self.inclination)
        z_ecf = y_ecf1 * np.sin(self.inclination) + z_ecf1 * np.cos(self.inclination)
        '''
        x_ecf = self.orbit_radius * (np.cos(self.raan) * np.cos(true_anomaly) - np.sin(self.raan) * np.sin(true_anomaly) * np.cos(self.inclination))
        y_ecf = self.orbit_radius * (np.sin(self.raan) * np.cos(true_anomaly) + np.cos(self.raan) * np.sin(true_anomaly) * np.cos(self.inclination))
        z_ecf = self.orbit_radius * np.sin(true_anomaly) * np.sin(self.inclination)
        return (x_ecf, y_ecf, z_ecf)
    # 获取所有卫星的位置
    def get_all_satellite_positions(self, time: float) -> List[Tuple[float, float, float]]:     
        # 遍历所有卫星，获取每个卫星的位置
        return [self.get_satellite_position(i, time) for i in range(self.satellite_count)]
    
    def get_orbit_points(self, num_points: int = 100) -> List[Tuple[float, float, float]]:
        """计算轨道上的点（用于绘制轨道）"""
        true_anomalies = np.linspace(0, 2*np.pi, num_points)
        orbit_points = []
        
        for ta in true_anomalies:
            x = self.orbit_radius * (np.cos(self.raan) * np.cos(ta) - np.sin(self.raan) * np.sin(ta) * np.cos(self.inclination))
            y = self.orbit_radius * (np.sin(self.raan) * np.cos(ta) + np.cos(self.raan) * np.sin(ta) * np.cos(self.inclination))
            z = self.orbit_radius * np.sin(ta) * np.sin(self.inclination)
            orbit_points.append((x, y, z))
            
        return orbit_points
    
    def get_normal_vector(self) -> Tuple[float, float, float]:
        return (np.sin(self.inclination) * np.sin(self.raan), -np.sin(self.inclination) * np.cos(self.raan), np.cos(self.inclination))